#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
from time import sleep

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper


def writeIpv4Rules(p4info_helper, sw, match_fields, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": match_fields
        },
        action_name="MyIngress.ipv4_forward",
        action_params=action_params
    )
    sw.WriteTableEntry(table_entry)

def writeEcnRules(p4info_helper, sw, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyEgress.check_ecn",
        action_name="MyEgress.mark_ecn",
        action_params=action_params,
        default_action=True
    )
    sw.WriteTableEntry(table_entry)

def writeCloneRules(p4info_helper, sw, clone_session_id, replicas):
    SessionEntry = p4info_helper.buildCloneSessionEntry(
        clone_session_id=clone_session_id,
        replicas=replicas
    )
    sw.WritePREEntry(SessionEntry)

def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 创建所有三个交换机的连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # 用户输入ECN阈值
        threshold = input("Please input the ECN threshold (queue depth): ")
        threshold = int(threshold)

        # 建立所有交换机的主仲裁
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # 安装P4程序到所有交换机
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                      bmv2_json_file_path=bmv2_file_path)
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                      bmv2_json_file_path=bmv2_file_path)
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                      bmv2_json_file_path=bmv2_file_path)

        # 为所有交换机写入完整的路由规则
        # s1 的路由规则
        writeIpv4Rules(p4info_helper, s1, ["10.0.1.1", 32],
                      {"dstAddr": "08:00:00:00:01:01", "port": 2})
        writeIpv4Rules(p4info_helper, s1, ["10.0.1.11", 32],
                      {"dstAddr": "08:00:00:00:01:11", "port": 1})
        writeIpv4Rules(p4info_helper, s1, ["10.0.2.0", 24],
                      {"dstAddr": "08:00:00:00:02:00", "port": 3})
        writeIpv4Rules(p4info_helper, s1, ["10.0.3.0", 24],
                      {"dstAddr": "08:00:00:00:03:00", "port": 4})

        # s2 的路由规则
        writeIpv4Rules(p4info_helper, s2, ["10.0.2.2", 32],
                      {"dstAddr": "08:00:00:00:02:02", "port": 2})
        writeIpv4Rules(p4info_helper, s2, ["10.0.2.22", 32],
                      {"dstAddr": "08:00:00:00:02:22", "port": 1})
        writeIpv4Rules(p4info_helper, s2, ["10.0.1.0", 24],
                      {"dstAddr": "08:00:00:00:01:00", "port": 3})
        writeIpv4Rules(p4info_helper, s2, ["10.0.3.0", 24],
                      {"dstAddr": "08:00:00:00:03:00", "port": 4})

        # s3 的路由规则
        writeIpv4Rules(p4info_helper, s3, ["10.0.3.3", 32],
                      {"dstAddr": "08:00:00:00:03:03", "port": 1})
        writeIpv4Rules(p4info_helper, s3, ["10.0.1.0", 24],
                      {"dstAddr": "08:00:00:00:01:00", "port": 2})
        writeIpv4Rules(p4info_helper, s3, ["10.0.2.0", 24],
                      {"dstAddr": "08:00:00:00:02:00", "port": 3})

        # 为所有交换机配置ECN规则
        writeEcnRules(p4info_helper, s1, {"ecn_threshold": threshold})
        writeEcnRules(p4info_helper, s2, {"ecn_threshold": threshold})
        writeEcnRules(p4info_helper, s3, {"ecn_threshold": threshold})

        # 为所有交换机配置克隆会话
        writeCloneRules(p4info_helper, s1, 100, [{"egress_port": 255, "instance": 1}])
        writeCloneRules(p4info_helper, s2, 100, [{"egress_port": 255, "instance": 1}])
        writeCloneRules(p4info_helper, s3, 100, [{"egress_port": 255, "instance": 1}])

        print(f"ECN threshold {threshold} has been configured successfully on all switches!")
        print("Controller is running... Press Ctrl+C to stop.")

        # 保持运行
        while True:
            sleep(10)

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/ecn.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/ecn.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
